
"""
Multimodal Data Loader with Text Features
Handles joint loading of rating data and text embeddings
"""

import torch
from torch.utils.data import Dataset

class MultimodalDataset(Dataset):
    def __init__(self, ratings, text_features):
        """
        Args:
            ratings: DataFrame of user-item ratings
            text_features: Dictionary of text embeddings 
                          {user_id: tensor, item_id: tensor}
        """
        self.ratings = ratings
        self.text_features = text_features

    def __len__(self):
        return len(self.ratings)

    def __getitem__(self, idx):
        user, item, rating = self.ratings.iloc[idx]
        return {
            'user': user,
            'item': item,
            'rating': torch.FloatTensor([rating]),
            'user_text': self.text_features['users'][user],
            'item_text': self.text_features['items'][item]
        }
